import cStringIO
import math
import numpy


class Generator:
	'''
	Generates a Neuron .hoc file to run a simulation according to the parameters specified. 
	'''
	
	def __init__(self, numPyrs, numInhib, connectivity, dt, duration, voltFilePre, voltFilePost):
		'''
		@param numPyrs: The number of excitatory cells (pyramidals)
		@param numInhib: The number of inhibitory cells
		@param connectivity: square matrix of dimension (numPyr+numInhib), with the pyrs before the inhibs. zero means no synapse. a 1 in [i,j] means a synapse from i to j.
		@param dt: timestep, in milliseconds
		@param duration: duration of the simulation , in milliseconds
		@param voltFilePre: each simulated cell outputs its voltage trace to a file vfPreNvfPost where N is the number of the cell (pyrs start at 0, inhibs start at numPyrs)
		@param voltFilePost: each simulated cell outputs its voltage trace to a file vfPreNvfPost where N is the number of the cell (pyrs start at 0, inhibs start at numPyrs)
		'''
		
		
		self.numPyrs = numPyrs
		self.numInhib = numInhib
		self.numCells = numPyrs+numInhib
		self.connectivity = connectivity
		self.dt = dt
		self.duration = duration
		self.voltageFilePre=voltFilePre
		self.voltageFilePost=voltFilePost
		self.numTimePoints=int(math.ceil(self.duration / self.dt))
		
		self.contents = cStringIO.StringIO()	
		self.cellString = "cells"
		self.synString = "synapses"
		self.synInd = 0
		
		
		
		
	#this is a function for making synapses from one thing to another thing!
	#it appends to the outstring.
	def makeSynapse(self, preInd, postInd, tau1, tau2, reversal):
		''' Sets up a synapse (alpha function) in the Neuron hoc file
		
		@param preInd: index (in the Neuron array referred to by self.cellString) of the presynaptic cell
		@param postInd: index of the postsynaptic cell
		@param tau1:  rising time constant 
		@param tau2:  decaying time constant
		@param reversal: reversal potential of the synapse
		'''
		
		#you make a synapse on the postsynaptic, and then use netcon to hook it up to the presynaptic.
		print >>self.contents, "access %s[%s].soma" % (self.cellString, postInd)
		print >>self.contents, 'synapses[%s] = new Exp2Syn(0.5)' % self.synInd #dual exponential synapse. the argument specifies location in the currently active compartment.
		print >>self.contents, 'synapses[%s].tau1 = %s' % (self.synInd, tau1)  #rise time constant (milliseconds?)
		print >>self.contents, 'synapses[%s].tau2 = %s' % (self.synInd, tau2)   #decay time constant (milliseconds?)
		print >>self.contents, 'synapses[%s].e = %s' % (self.synInd, reversal)  #reversal potential (mV)
		#netcon-it : 
		#//netcon = new NetCon(source, target, threshold, delay, weight)
		#first access the pre-synaptic compartment
		print >>self.contents, 'access %s[%s].soma' % (self.cellString, preInd)
		#then do it:
		print >>self.contents, '%s[%s].nclist.append(new NetCon(&v(0.5), synapses[%s], -20, 1, 0.004))' % (self.cellString, postInd, self.synInd)
		self.synInd +=1
	
		return

	def generate(self):
		''' Creates the text for the Neuron .hoc file, accessable as a CStringIO.StringIO object via this.contents
		'''
		#neuron-setup-stuff  .. mandatory
		print >>self.contents, 'load_file(\"pyrTemplate.hoc\")'
		print >>self.contents, 'load_file(\"nrngui.hoc\")'
		
		#basic simulation variables
		print >>self.contents, "dt = %s" % self.dt
		print >>self.contents, "npoints = %s" % self.numTimePoints
		print >>self.contents, "tstop = %s" % self.duration
		print >>self.contents, 'steps_per_ms = 1/dt'
		print >>self.contents, 'v_init = -67'
		print >>self.contents, 'celsius = 36'
		
		
		###################network setup#################
		
		#setup cells, and their voltage vectors.
		print >>self.contents, "objectvar %s[%s]" % (self.cellString, self.numCells)
		print >>self.contents, "objectvar volts[%s]" % self.numCells
		print >>self.contents, "for i=0, %s{" % int(self.numCells -1)
		print >>self.contents, "	%s[i] = new MacPyr()" % self.cellString
		print >>self.contents, "	volts[i] = new Vector()"
		print >>self.contents, "	volts[i].record(&%s[i].soma.v(0.5))" % self.cellString
		print >>self.contents, "}"
		
		#setup synapses
		#make an array:
		
		#connections = numpy.zeros(numCells**2, dtype=int).reshape(numCells,numCells)
		#connections[0,1] = 1
		#connections[0,2] = 1
		#connections[2,1] = 1
		#print connections
		
		#setup the synapses in Neuron
		#the synapses must all be in an array
		print >>self.contents, "objectvar synapses[%s]" % numpy.sum(self.connectivity)
		
		#iterate across the connectivity matrix:
		#in slot i,j  a 0 means no synapse, and a 1 means a synapse from i to j.
		#the first numPyrs cells are pyramidals, the remainder are inhibitory.
		for i in range(0, self.numCells):
			for j in range(0, self.numCells):
				if(self.connectivity[i,j] == 0):
					continue
				revPot = 0 #the reversal potential of the synapse
				if(i >= self.numPyrs):
					revPot = -80	
				print "calling makeSynapse at %s,%s , revpot=%s" %(i, j, revPot)
				
				self.makeSynapse(i, j, .1, 10, revPot)
			
		
		#setup stimulating electrode:
		#//we want to stimulate one cell , so:
		#//the 'active' cell is the one we mention
		#//most recently,  so we just say 
		#// 'the zeroth pyramidal' ,  and that's
		#//where the electrode goes:
		print >>self.contents, 'objectvar macStim'
		print >>self.contents, "access %s[0].soma" % self.cellString
		print >>self.contents, 'macStim = new IClamp(0.5)' #the somas are currently points so placing at 0.5 means "the middle" but also "the only point"
		print >>self.contents, 'macStim.del = 30' #start the stim at a 50 millisecond delay from time zero
		print >>self.contents, 'macStim.dur = 150' #duration of stim in milliseconds
		print >>self.contents, 'macStim.amp = 0.35' #amplitude of stim in nanoAmps
		
		
		#record the timesteps so we can output them and use them to plot our data..
		#there must be some current segment selected, even though time is global.
		print >>self.contents, "objref times"
		print >>self.contents, "times = new Vector()"
		print >>self.contents, "times.record(&t)"
		
		
		###### RUN  THE SIM
		
		#the commands to run the sim must precede the commands to save the simulation data ? i think.
		print >>self.contents, 'init()'
		print >>self.contents, 'run()'
		
		
		###
		
		
		#setup the simulation's file output:
		#simulation times:
		print >>self.contents, "objref saver"
		print >>self.contents, "saver = new File()"
		print >>self.contents, "saver.wopen(\"timesteps.txt\")"
		print >>self.contents, "times.printf(saver)"
		print >>self.contents, "saver.close()"
		
		
		#voltage from each pyramidal cell.
		for i in range(0, self.numCells):
			print >>self.contents, "saver = new File()"
			print >>self.contents, "saver.wopen(\"%s%s%s\")" % (self.voltageFilePre, i, self.voltageFilePost)
			print >>self.contents, "volts[%s].printf(saver)" % i
			print >>self.contents, "saver.close()"
		
	def writeFile(self, filename):
		''' 
		after you call self.generate, you can call this to write self.contents to a file
		@param filename: file name
		'''
		#open the file for writing
		mainfile = open(filename, 'w')
		try:
		    mainfile.write(self.contents.getvalue())
		#make sure the file gets closed.
		finally:
		    mainfile.close()
		#free the string on the way out
		#self.contents.close()

	